from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import backtrader as bt
import os
import datetime
import pandas as pd

# --- 假设之前的策略、Sizer、数据加载器等都定义在 backtrader_strategy.py 中 ---
try:
    from backtrader_strategy import SystematicRiskParityStrategyBT, RiskParitySizer, StampDutyCommissionScheme, GenericCSV_X
except ImportError:
    print("错误：无法导入 'backtrader_strategy.py'。")
    print("请确保此优化脚本与 'backtrader_strategy.py' 在同一个文件夹下。")
    exit()


if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # --- 1. 使用 optstrategy 来添加策略 ---
    cerebro.optstrategy(
        SystematicRiskParityStrategyBT,
        sma_long_window=range(140, 161, 10) # 测试 140, 150, 160 三个值
    )

    # --- 2. 加载数据 ---
    data_path = './data'
    universe = [
        "000651.SZ",
        "600000.SH",
        "600036.SH",
        "600519.SH",
    ]
    print(f"Loading {len(universe)} data feeds for optimization...")
    for stock in universe:
        dataname = os.path.join(data_path, f'{stock}.csv')
        if os.path.exists(dataname):
            feed = GenericCSV_X(dataname=dataname, fromdate=datetime.datetime(2020, 1, 1), todate=datetime.datetime(2024, 12, 31))
            cerebro.adddata(feed, name=stock)

    # --- 3. 配置环境 ---
    cerebro.addsizer(RiskParitySizer)
    
    # 添加所有需要的分析器
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharperatio')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')
    cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='tradeanalyzer')
    # ######################################################
    # ##                核心修改在这里 (1/2)                 ##
    # ######################################################
    # 添加 Returns 分析器，用于在优化中获取最终回报率
    cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')
    # ######################################################
    
    initial_cash = 1000000.0
    cerebro.broker.setcash(initial_cash)
    comm_scheme = StampDutyCommissionScheme()
    cerebro.broker.addcommissioninfo(comm_scheme)
    
    # --- 4. 运行优化 ---
    print("\n--- Starting Parameter Optimization (Single Core Mode) ---")
    # 明确禁用多进程并行计算，确保稳定运行
    opt_results = cerebro.run(maxcpus=1)

    # --- 5. 解析并打印优化结果 ---
    print("\n--- Optimization Results ---")
    final_results = []
    for run in opt_results:
        for strategy_result in run:
            params = strategy_result.p._getkwargs()
            
            # ######################################################
            # ##                核心修改在这里 (2/2)                 ##
            # ######################################################
            # 从 Returns 分析器中获取总回报率，并计算最终价值
            returns_analysis = strategy_result.analyzers.returns.get_analysis()
            total_return = returns_analysis.get('rtot', 0.0)
            final_value = initial_cash * (1 + total_return)
            # ######################################################

            sharpe = strategy_result.analyzers.sharperatio.get_analysis().get('sharperatio', 0.0)
            drawdown = strategy_result.analyzers.drawdown.get_analysis().get('max', {}).get('drawdown', 0.0)
            trade_info = strategy_result.analyzers.tradeanalyzer.get_analysis()
            total_trades = trade_info.get('total', {}).get('total', 0)
            
            final_results.append({
                'sma_long_window': params['sma_long_window'],
                'final_value': final_value,
                'sharpe_ratio': sharpe,
                'max_drawdown': drawdown,
                'total_trades': total_trades
            })
    
    # 使用 Pandas DataFrame 进行格式化和排序
    results_df = pd.DataFrame(final_results)
    results_df.sort_values(by='final_value', ascending=False, inplace=True)

    # 格式化输出
    results_df['final_value'] = results_df['final_value'].map('{:,.2f}'.format)
    results_df['sharpe_ratio'] = results_df['sharpe_ratio'].map('{:.3f}'.format)
    results_df['max_drawdown'] = results_df['max_drawdown'].map('{:.2%}'.format)
    
    print(results_df.to_string())